# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Use the distances between atoms to calculate neighbour list
"""

from typing import Tuple
import torch
from torch import Tensor
from torch.nn import Parameter, Module

from ..function.functions import get_integer, get_tensor, reduce_all
from ..function.operations import GetDistance


class DistanceNeighbours(Module):
    r"""Neighbour list calculated by distance

    Args:

        cutoff (float):         Cutoff distance.

        num_neighbours (int):   Number of neighbours. If `None` is given, this value will be calculated by
                                the ratio of the number of neighbouring grids to the total number of grids.
                                Default: None

        atom_mask (Tensor):     Tensor of shape `(B, A)`. Data type is bool.
                                Mask of atoms in the system.
                                Default: None

        exclude_index (Tensor): Tensor of shape `(B, A, Ex)`. Data type is int32.
                                Index of neighbour atoms which could be excluded from the neighbour list.
                                Default: None

        use_pbc (bool):         Whether to use periodic boundary condition. Default: None

        cutoff_scale (float):   Factor to scale the cutoff distance. Default: 1.2

        large_dis (float):      A large number to fill in the distances to the masked neighbouring atoms.
                                Default: 1e4

        cast_fp16 (bool):       If this is set to `True`, the data will be cast to float16 before sort.
                                For use with some devices that only support sorting of float16 data.
                                Default: False

    Supported Platforms:

        ``Ascend`` ``GPU``

    Symbols:

        B:  Number of simulation walker.

        A:  Number of atoms in system.

        N:  Number of neighbour atoms.

        Ex: Maximum number of excluded neighbour atoms.

    """

    def __init__(self,
                 cutoff: float = None,
                 num_neighbours: int = None,
                 atom_mask: Tensor = None,
                 exclude_index: Tensor = None,
                 use_pbc: bool = None,
                 cutoff_scale: float = 1.2,
                 large_dis: float = 1e4,
                 ):

        super().__init__()

        self.cutoff = get_tensor(cutoff, torch.float32)
        self.cutoff_scale = get_tensor(cutoff_scale, torch.float32)
        self.scaled_cutoff = self.cutoff * self.cutoff_scale

        self.num_neighbours = get_integer(num_neighbours)
        if self.num_neighbours is None:
            max_neighbours = get_tensor(0, dtype=torch.int32)
        else:
            max_neighbours = get_tensor(self.num_neighbours, dtype=torch.int32)
        self.max_neighbours = max_neighbours

        self.large_dis = get_tensor(large_dis, dtype=torch.float32)

        self.emtpy_atom_shift = 0
        self.atom_mask = None
        self.has_empty_atom = False
        if atom_mask is not None:
            # (B, A)
            self.atom_mask = get_tensor(atom_mask, dtype=torch.bool)
            if self.atom_mask.ndim == 1:
                self.atom_mask = self.atom_mask.unsqueeze(0)

            self.has_empty_atom = not self.atom_mask.all()
            if self.has_empty_atom:
                emtpy_atom_mask = ~self.atom_mask
                # (B, 1, A)
                self.emtpy_atom_shift = emtpy_atom_mask.unsqueeze(-2) * self.large_dis

        self.exclude_index = None
        if exclude_index is not None:
            # (B, A, Ex)
            self.exclude_index = get_tensor(exclude_index, dtype=torch.int32)
            if self.exclude_index.ndim == 2:
                self.exclude_index = self.exclude_index.unsqueeze(0)

        self.get_distance = GetDistance(use_pbc)

    def calc_max_neighbours(self, distances: Tensor, cutoff: float) -> int:
        """计算每个原子的最大邻居数

        Args:
            distances (Tensor): 形状为 (B, A, A) 的距离张量
            cutoff (float): 截断距离

        Returns:
            int: 最大邻居数
        """
        neighbour_mask = distances < cutoff
        num_neighbours = torch.count_nonzero(neighbour_mask, dim=-1)
        return torch.max(num_neighbours)

    def set_num_neighbours(self,
                           coordinate: Tensor,
                           pbc_box: Tensor = None,
                           scale_factor: float = 1.25
                           ):
        """set maximum number of neighbouring atoms"""
        distances = self.get_distance(coordinate.unsqueeze(-2), coordinate.unsqueeze(-3), pbc_box)
        num_neighbours = self.calc_max_neighbours(distances, self.scaled_cutoff)
        num_neighbours = torch.ceil(num_neighbours * scale_factor)
        self.num_neighbours = get_integer(torch.minimum(num_neighbours, get_tensor(coordinate.shape[-2] - 1)))
        self.max_neighbours = self.num_neighbours
        return self

    def set_exclude_index(self, exclude_index: Tensor) -> Tensor:
        # (B, A, Ex)
        self.exclude_index = get_tensor(exclude_index, dtype=torch.int32)
        if self.exclude_index.ndim == 2:
            self.exclude_index = self.exclude_index.unsqueeze(0)
        return self.exclude_index

    def print_info(self):
        print('[MindSPONGE] Neighbour list: DistanceNeighbours')
        print(f'[MindSPONGE]     Cut-off distance: {self.cutoff}')
        print(f'[MindSPONGE]     Scaled cut-off: {self.scaled_cutoff}')
        print(f'[MindSPONGE]     Max number of neighbour atoms: {self.num_neighbours}')
        return self

    def check_neighbour_list(self, raise_error: bool = True) -> bool:
        """check the number of neighbouring atoms in neighbour list"""
        # if self.num_neighbours is not None and self.max_neighbours > self.num_neighbours:
        #     if raise_error:
        #         raise RuntimeError(f'The max number of neighbour atoms ({self.max_neighbours.numpy()}) '
        #                            f'is larger than the initial neighbouring number of neighbour list '
        #                            f'({self.num_neighbours}!')
        #     print('[WARNING] The max number of neighbour atoms exceeds '
        #           'the initial neighbouring number of neighbour list!')
        #     return False
        return True

    def forward(self,
                  coordinate: Tensor,
                  pbc_box: Tensor = None,
                  atom_mask: Tensor = None,
                  exclude_index: Tensor = None
                  ) -> Tuple[Tensor, Tensor]:
        """计算邻居列表

        Args:
            coordinate (Tensor): 形状为 (B, A, D) 的坐标张量
            pbc_box (Tensor): 形状为 (B, D) 的周期性边界条件盒子
            atom_mask (Tensor): 形状为 (B, A) 的原子掩码
            exclude_index (Tensor): 形状为 (B, A, Ex) 的排除索引

        Returns:
            Tuple[Tensor, Tensor]: 邻居索引和邻居掩码
        """
        # 计算距离矩阵
        distances = self.get_distance(coordinate.unsqueeze(-2), coordinate.unsqueeze(-3), pbc_box)

        # 处理原子掩码
        if atom_mask is None:
            atom_mask = self.atom_mask
            if self.has_empty_atom:
                distances += self.emtpy_atom_shift
        else:
            if not atom_mask.all():
                emtpy_atom_mask = ~atom_mask
                emtpy_atom_shift = emtpy_atom_mask.unsqueeze(-2) * self.large_dis
                distances += emtpy_atom_shift

        # 确定邻居数量
        num_atoms = coordinate.shape[-2]
        if self.num_neighbours is None:
            num_neighbours = num_atoms - 1
        else:
            num_neighbours = self.num_neighbours
            max_neighbours = self.calc_max_neighbours(distances, self.cutoff)
            self.max_neighbours = max_neighbours

        distances, neighbours = torch.topk(-distances, num_neighbours+1, dim=-1)
        distances = -distances[..., 1:]
        neighbours = neighbours[..., 1:]
        neighbour_mask = distances < self.scaled_cutoff

        # 处理排除索引
        if exclude_index is None:
            exclude_index = self.exclude_index

        if exclude_index is not None:
            exclude_mask = neighbours.unsqueeze(-1) != exclude_index.unsqueeze(-2)
            exclude_mask =  reduce_all(exclude_mask, -1)
            neighbour_mask = torch.logical_and(exclude_mask, neighbour_mask)

        if atom_mask is not None:
            neighbour_mask = torch.logical_and(neighbour_mask, atom_mask.unsqueeze(-1))

        # (1, A, 1)
        no_idx = torch.arange(num_atoms).reshape(1, -1, 1)
        # (B, A, n)
        no_idx = torch.broadcast_to(no_idx, neighbours.shape)
        neighbours = torch.where(neighbour_mask, neighbours, no_idx)        

        return distances, neighbours, neighbour_mask
